# ChatDoc:和文件聊天
from langchain_community.document_loaders import UnstructuredExcelLoader, Docx2txtLoader, PyPDFLoader
from langchain_text_splitters import CharacterTextSplitter
from langchain_chroma import Chroma
from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain_openai import OpenAIEmbeddings
from langchain_openai import ChatOpenAI
from langchain_deepseek import ChatDeepSeek
from langchain_core.prompts import ChatPromptTemplate
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import LLMChainExtractor
import os

apiproxy_api_key = "sk-tjodGSqf2LSLFnn7YTnwJUZkFfIMMEBiNeOAquBNJPCp6jSI"
apiproxy_base_url = "https://apiproxy.top/v1"
embeddings = OpenAIEmbeddings(
    model="text-embedding-3-small",
    api_key=apiproxy_api_key,
    base_url=apiproxy_base_url
)

os.environ["DASHSCOPE_API_BASE"] = 'https://api.deepseek.com'
os.environ["DEEPSEEK_API_KEY"] = 'sk-66899d5ed8094f91904f73424aa380aa'
llm = ChatDeepSeek(
    model="deepseek-reasoner",
    temperature=0,
)

# llm_openai = ChatOpenAI(
#     model="gpt-4o-mini",
#     api_key=apiproxy_api_key,
#     base_url=apiproxy_base_url
# )


class ChatDoc():
    def __init__(self):
        self.doc = None
        self.splitText = []
        self.template = [
            ("system", "你是一个处理文档的秘书，你从不说自己是一个大模型或者AI助手，你会根据下面提供的上下文内容来继续回答问题.\n上下文内容\n{context}\n"),
            ("human", "你好"),
            ("ai", "你好"),
            ("human", "{question}"),
        ]
        self.prompt = ChatPromptTemplate.from_messages(self.template)

    #加载文档
    def getFile(self):
        doc = self.doc
        loaders = {
            "docx": Docx2txtLoader,
            "pdf": PyPDFLoader,
            "xlsx": UnstructuredExcelLoader
        }
        file_extension = doc.split(".")[-1]
        loader_class = loaders[file_extension]
        if loader_class:
            try:
                loader = loader_class(doc)
                text = loader.load()
                return text
            except Exception as e:
                print(f"Error loading {file_extension} files: {e}")
        else:
            print(f"Unsupported file extension: {file_extension}")
            return None

    #处理文档
    def splitSentence(self):
        full_text = self.getFile()
        if full_text != None:
            #对文档进行分割
            text_split = CharacterTextSplitter(
                chunk_size=150,
                chunk_overlap=20
            )
            texts = text_split.split_documents(full_text)
            self.splitText = texts

    def embeddingAndVectorDB(self):
        # print(self.splitText)
        db = Chroma.from_documents(
            documents=self.splitText,
            embedding=embeddings
        )
        return db

    #提问并找到相关的文本块
    def askAndFindFiles(self, question):
        db = self.embeddingAndVectorDB()
        retriever = db.as_retriever(search_type="mmr")
        results = retriever.invoke(input=question)
        return results

    def chatWithDoc(self, question):
        _content = ""
        context = self.askAndFindFiles(question)
        for i in context:
            _content += i.page_content

        messages = self.prompt.format_messages(context=_content, question=question)
        return llm.invoke(messages)


